#include "login_client_imp.h"

#include "hex_dumper.h"

namespace login_client_detail
{
	bool ChallengePhase::handle(IPackage* package, Imp* loginClient, IIoHandler::Ptr ioHandler)
	{
		if(package->messageID() != ChallengePackage::messageID_)
		{
			return false;
		}
		auto challengePackage = static_cast<ChallengePackage*>(package);
		if(!challengePackage)
		{
			return false;
		}
		HMAC<SHA> hmac(loginClient->digestOfUserPassword_, SHA::DIGESTSIZE);

		{
			cout << "challenge phase, digest of user pass: " << 
				hex_dump::dump(loginClient->digestOfUserPassword_, SHA::DIGESTSIZE) << endl;
		}

		auto response = new ChallengeResponse();
		response->nonce_.reset(new unsigned char[SHA::DIGESTSIZE], SHA::DIGESTSIZE);
		response->userName_ = loginClient->userName_;

		hmac.CalculateDigest((byte*)response->nonce_.get(), 
			(const byte*)challengePackage->nonce_.get(), challengePackage->nonce_.size());

		{	
			cout << "challenge phase, hmac of server nones: " << 
				hex_dump::dump(response->nonce_.get(), SHA::DIGESTSIZE) << endl;
		}

		ioHandler->send(IPackage::Ptr(response));

		loginClient->currentState_.reset(new KeyExchangePhase()); 

		return true;
	}


	bool KeyExchangePhase::handle(IPackage* package, Imp* loginClient, IIoHandler::Ptr ioHandler)
	{
		if(package->messageID() != KeyExchangePackage::messageID_)
		{
			return false;
		}
		auto keyExchangePackage = static_cast<KeyExchangePackage*>(package);
		HMAC<SHA> hmac(loginClient->digestOfUserPassword_, SHA::DIGESTSIZE);
		hmac.Update((const byte*)loginClient->userName_.data(), loginClient->userName_.length());
		hmac.Update((const byte*)loginClient->password_.data(), loginClient->password_.length());
		hmac.Update((const byte*)keyExchangePackage->nonce_.get(), 
			keyExchangePackage->nonce_.size());

		auto inKey = scoped_array<byte>(new byte[SHA::DIGESTSIZE], SHA::DIGESTSIZE);
		hmac.Final(inKey.get());

		auto inCodec = new Arc4Codec(inKey);
		{
			cout << "key exchange phase, in key: " << 
				hex_dump::dump(inKey.get(), SHA::DIGESTSIZE) << endl;
		}

		AutoSeededRandomPool rng;
		auto response = new KeyExchangeResponse();
		response->nonce_.reset(new byte[SHA::DIGESTSIZE], SHA::DIGESTSIZE);
		rng.GenerateBlock((byte*)response->nonce_.get(), SHA::DIGESTSIZE);

		hmac.Restart();
		hmac.Update((const byte*)loginClient->userName_.data(), loginClient->userName_.length());
		hmac.Update((const byte*)loginClient->password_.data(), loginClient->password_.length());
		hmac.Update((const byte*)response->nonce_.get(), SHA::DIGESTSIZE);

		scoped_array<byte> outKey(new byte[SHA::DIGESTSIZE], SHA::DIGESTSIZE);
		hmac.Final((byte*)outKey.get());

		ioHandler->send(IPackage::Ptr(response));

		auto outCodec = new Arc4Codec(outKey);
		{
			cout << "key exchange phase, out key: " << 
				hex_dump::dump(outKey.get(), SHA::DIGESTSIZE) << endl;
		}

		ioHandler->setAttachment(
			IIoAttachment::Ptr(new DuplexCodec(IPackageCodec::Ptr(outCodec), 
			IPackageCodec::Ptr(inCodec))));

		if(loginClient->loginSuccessed_)
			loginClient->loginSuccessed_(ioHandler);

		return true;
	}


	Imp::Imp(const string& userName, const string& password) : 
	userName_(userName), password_(password), currentState_(new ChallengePhase())
	{
		SHA sha;
		sha.Update((const byte*)userName.data(), userName.length());
		sha.Update((const byte*)password.data(), password.length());
		sha.Final(digestOfUserPassword_);
	}


	void Imp::handle(IPackage* package, IIoHandler::Ptr ioHandler)
	{
		if(!currentState_->handle(package, this, ioHandler))
		{
			if(loginFailed_)
				loginFailed_(ioHandler);
			ioHandler->close();
		}
	}
}

using namespace login_client_detail;
LoginClient::LoginClient(const string& userName, const string& password) :
	imp_(new login_client_detail::Imp(userName, password))
{

}


void LoginClient::handle(IPackage* package, IIoHandler::Ptr ioHandler)
{
	imp_->handle(package, ioHandler);
}

void LoginClient::setLoginFailedDelegate(LoginAuthDelegate loginFailed)
{
	imp_->setLoginFailedDelegate(loginFailed);
}
void LoginClient::setLoginSuccessDelegate(LoginAuthDelegate loginSuccessed)
{
	imp_->setLoginSuccessDelegate(loginSuccessed);
}